// Socket server that can be shut down -- stop serving, in a graceful manner.
// This version shuts down idle open connections from clients.
//
// Eli Bendersky [https://eli.thegreenplace.net]
// This code is in the public domain.

// https://github.com/eliben/code-for-blog/blob/main/2020/tcp-server-shutdown/shutdown2/shutdown2.go

package main

import (
	"io"
	"log"
	"net"
	"sync"
	"time"
)

type TcpServerHandler interface {
	OnRecv(conn net.Conn, msg []byte)
}

type TcpServer struct {
	listener net.Listener
	quit     chan interface{}
	wg       sync.WaitGroup
	handler  TcpServerHandler
}

func NewTcpServer(addr string, handler TcpServerHandler) *TcpServer {
	s := &TcpServer{
		listener: nil,
		quit:     make(chan interface{}),
		handler:  handler,
	}

	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		log.Fatal(err)
	}

	l, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		log.Fatal(err)
	}

	s.listener = l
	s.wg.Add(1)
	go s.serve()
	return s
}

func (s *TcpServer) Stop() {
	close(s.quit)
	s.listener.Close()
	s.wg.Wait()
}

func (s *TcpServer) Addr() string {
	return s.listener.Addr().String()
}

func (s *TcpServer) Network() string {
	return s.listener.Addr().Network()
}

func (s *TcpServer) serve() {
	defer s.wg.Done()

	for {
		conn, err := s.listener.Accept()
		if err != nil {
			select {
			case <-s.quit:
				log.Printf("Closing server %s %s", s.Network(), s.Addr())
				return
			default:
				log.Println("Accept error", err)
			}
		} else {
			s.wg.Add(1)
			log.Printf("Accepted %v", conn.RemoteAddr())

			go func() {
				s.handleConection(conn)
				s.wg.Done()
			}()
		}
	}
}

func (s *TcpServer) handleConection(conn net.Conn) {
	defer conn.Close()
	if nil == s.handler {
		log.Printf("Handler is nil, force close %v", conn.RemoteAddr())
		return
	}

	buf := make([]byte, 2048)
ReadLoop:
	for {
		select {
		case <-s.quit:
			return
		default:
			conn.SetDeadline(time.Now().Add(200 * time.Millisecond))
			n, err := conn.Read(buf)
			if err != nil {
				if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
					continue ReadLoop
				} else if err != io.EOF {
					log.Println("Read error", err)
					return
				}
			}
			if n == 0 {
				log.Printf("Closed by %v", conn.RemoteAddr())
				return
			}

			s.handler.OnRecv(conn, buf[:n])
		}
	}
}

/*
func init() {
	log.SetFlags(log.Ltime | log.Lmicroseconds)
}
*/
